﻿// Disk: Disk and hardware module
// Copyright (C) 2014 Nathan Moschkin.

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Runtime.InteropServices;
using DataTools.Memory;
using DataTools;
using DataTools.Interop;

namespace DataTools.Interop.Hardware
{
    // LONG STORY SHORT:
    // -------------------------------------------------------------------------------------------
    // CSDisk.DiskInfo[] disks = CSDisk.DiskUtil.EnumDisks();


    #region Disk Structures

    /// <summary>
    /// Disk Media Types
    /// </summary>
    public enum MediaType
    {
        /// <summary>
        /// Unknown Media.
        /// </summary>
        Unknown = 0x0,

        /// <summary>
        /// 
        /// </summary>
        F5_1Pt2_512 = 0x1,

        /// <summary>
        /// 
        /// </summary>
        F3_1Pt44_512 = 0x2,

        /// <summary>
        /// 
        /// </summary>
        F3_2Pt88_512 = 0x3,

        /// <summary>
        /// 
        /// </summary>
        F3_20Pt8_512 = 0x4,

        /// <summary>
        /// 
        /// </summary>
        F3_720_512 = 0x5,

        /// <summary>
        /// 
        /// </summary>
        F5_360_512 = 0x6,

        /// <summary>
        /// 
        /// </summary>
        F5_320_512 = 0x7,

        /// <summary>
        /// 
        /// </summary>
        F5_320_1024 = 0x8,

        /// <summary>
        /// 
        /// </summary>
        F5_180_512 = 0x9,

        /// <summary>
        /// 
        /// </summary>
        F5_160_512 = 0xA,

        /// <summary>
        /// Removable Media.
        /// </summary>
        RemovableMedia = 0xB,

        /// <summary>
        /// Fixed Media.
        /// </summary>
        FixedMedia = 0xC,

        /// <summary>
        /// 
        /// </summary>
        F3_120M_512 = 0xD,

        /// <summary>
        /// 
        /// </summary>
        F3_640_512 = 0xE,

        /// <summary>
        /// 
        /// </summary>
        F5_640_512 = 0xF,

        /// <summary>
        /// 
        /// </summary>
        F5_720_512 = 0x10,

        /// <summary>
        /// 
        /// </summary>
        F3_1Pt2_512 = 0x11,

        /// <summary>
        /// 
        /// </summary>
        F3_1Pt23_1024 = 0x12,

        /// <summary>
        /// 
        /// </summary>
        F5_1Pt23_1024 = 0x13,

        /// <summary>
        /// 
        /// </summary>
        F3_128Mb_512 = 0x14,

        /// <summary>
        /// 
        /// </summary>
        F3_230Mb_512 = 0x15,

        /// <summary>
        /// 
        /// </summary>
        F8_256_128 = 0x16,

        /// <summary>
        /// 
        /// </summary>
        F3_200Mb_512 = 0x17,

        /// <summary>
        /// 
        /// </summary>
        F3_240M_512 = 0x18,

        /// <summary>
        /// 
        /// </summary>
        F3_32M_512 = 0x19
    };

    /// <summary>
    /// Describes the disk geometry of a physical drive.
    /// </summary>
    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode), TypeConverter(typeof(ExpandableObjectConverter))]
    internal struct DISK_GEOMETRY
    {
        public long Cylinders;
        public MediaType MediaType;
        public uint TracksPerCylinder;
        public uint SectorsPerTrack;
        public uint BytesPerSector;

        public long TotalSectors
        {
            get
            {
                return (long)SectorsPerTrack * (long)TracksPerCylinder * (long)Cylinders;
            }
        }

        public override string ToString()
        {
            return Text.PrintFriendlySize((long)TotalSectors * (long)BytesPerSector) + ", Total Sectors: " + TotalSectors.ToString("0,0") + ", Bytes Per Sector: " + BytesPerSector.ToString("0,0");
        }
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode), TypeConverter(typeof(ExpandableObjectConverter))]
    internal struct DISK_GEOMETRY_EX
    {
        public DISK_GEOMETRY Geometry;
        public ulong DiskSize;
        public byte Data;

        public override string ToString()
        {
            return Geometry.ToString();
        }
    }

    #endregion Interop

    /// <summary>
    /// DiskUtil physical drive enumerator.
    /// </summary>
    public static class DiskUtil
    {
        /// <summary>
        /// Enumerates all visible disk drives on the local system.
        /// </summary>
        /// <returns></returns>
        public static DiskInfo[] EnumDisks()
        {
            List<DiskInfo> disks = new List<DiskInfo>();

            MemPtr mm = new MemPtr();
            IntPtr disk;

            uint c = 0;
            uint cbSize;

            // initialize a bunch of structures ...

            SP_DEVINFO_DATA devInfo = new SP_DEVINFO_DATA();
            SP_DEVICE_INTERFACE_DATA devInterface = new SP_DEVICE_INTERFACE_DATA();
            SP_DEVICE_INTERFACE_DETAIL_DATA devDetail = new SP_DEVICE_INTERFACE_DETAIL_DATA();
            STORAGE_DEVICE_NUMBER diskNumber;

            DiskInfo devOut;

            // We want disks (but you can do this with literally any kind of device)
            Guid ClassId = Setup.GUID_DEVINTERFACE_DISK;

            // Get the devices in the class (interface)!
            IntPtr hDev = Setup.SetupDiGetClassDevs(ClassId, IntPtr.Zero, IntPtr.Zero, ClassDevFlags.DeviceInterface | ClassDevFlags.Present);

            if (hDev == IoCtl.INVALID_HANDLE_VALUE)
            {
                return null;
            }

            // preset the cbSize of the structures so the OS knows what it's playing with...
            devInfo.cbSize = (uint)Marshal.SizeOf(devInfo);
            devInterface.cbSize = (uint)Marshal.SizeOf(devInterface);

            // and ... enumerate!
            while (Setup.SetupDiEnumDeviceInterfaces(hDev, IntPtr.Zero, ref ClassId, (uint)c, ref devInterface) != 0)
            {
                devOut = new DiskInfo();

                // get the property bag...

                Setup.SetupDiEnumDeviceInfo(hDev, c, ref devInfo);

                // we want these luxurious properties ...

                devOut.HardwareIds = (string[])Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_HardwareIds, Setup.DevPropTypes.StringList);
                devOut.FriendlyName = (string)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_FriendlyName, Setup.DevPropTypes.String);
                devOut.PDOName = (string)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_PDOName, Setup.DevPropTypes.String);
                devOut.InstallDate = (DateTime)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_FirstInstallDate, Setup.DevPropTypes.Date);
                devOut.DeviceInterfaceClassGuid = ClassId;
                devOut.DeviceClassGuid = (Guid)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_ClassGuid, Setup.DevPropTypes.Guid);
                devOut.Manufacturer = (string)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_Manufacturer, Setup.DevPropTypes.String);
                devOut.Description = (string)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_DeviceDesc, Setup.DevPropTypes.String);
                devOut.InstanceId = (string)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_InstanceId, Setup.DevPropTypes.String);
                devOut.LocationInfo = (string)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_LocationInfo, Setup.DevPropTypes.String);
                devOut.LocationPaths = (string[])Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_LocationPaths, Setup.DevPropTypes.String);
                devOut.Parent = (string)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_Parent, Setup.DevPropTypes.String);
                devOut.RemovalPolicy = (DeviceRemovalPolicy)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_RemovalPolicy, Setup.DevPropTypes.Int32);
                devOut.SafeRemovalRequired = (bool)Setup.GetDevProperty(hDev, devInfo, Setup.DEVPKEY_Device_SafeRemovalRequired, Setup.DevPropTypes.Boolean);

                cbSize = 0;

                // allocate our memory...
                mm.Alloc(Marshal.SizeOf(devDetail));

                // shortcut: not passing the whole structure ... just marking the size where it should be (position 0)...
                mm.intAt(0, Marshal.SizeOf(devDetail));

                // request the size by setting every input to null:
                Setup.SetupDiGetDeviceInterfaceDetail(hDev, ref devInterface, IntPtr.Zero, 0, ref cbSize, IntPtr.Zero);

                if (cbSize > 0)
                {
                    mm.ReAlloc(cbSize);

                    // now call for real ...
                    Setup.SetupDiGetDeviceInterfaceDetail(hDev,
                        ref devInterface,
                        mm,
                        cbSize,
                        ref cbSize,
                        IntPtr.Zero);

                    // we know this is here ...
                    devOut.DevicePath = mm.GrabString(4);

                    // that was the magic thing we needed!
                    disk = FileIO.CreateFile(devOut.DevicePath, FileIoAccess.Read, ShareModes.Read | ShareModes.Write, IntPtr.Zero, CreateDisposition.OpenExisting, 0, IntPtr.Zero);

                    if (disk != MemPtr.InvalidHandle)
                    {
                        // we have an open disk:
                        mm.ReAlloc(Marshal.SizeOf(new STORAGE_DEVICE_NUMBER()));
                        mm.ZeroMemory();

                        // get the storage device number (aka \\.\PhysicalDriveX)
                        // You can use either DevicePath or PhysicalDriveX to open a disk...
                        // but it takes one to get the other (without WMI)... and it's good to have both on hand...
                        IoCtl.DeviceIoControl(disk,
                                IoCtl.IOCTL_STORAGE_GET_DEVICE_NUMBER,
                                IntPtr.Zero,
                                0,
                                mm,
                                (uint)mm.GetLength(),
                                ref cbSize,
                                IntPtr.Zero);

                        if (cbSize != 0)
                        {
                            diskNumber = mm.ToStruct<STORAGE_DEVICE_NUMBER>();
                            devOut.PhysicalDevice = diskNumber.DeviceNumber;
                        }

                        mm.ReAlloc(Marshal.SizeOf(new DISK_GEOMETRY_EX()));

                        // Now, we get that disk geometry like you wanted...
                        IoCtl.DeviceIoControl(disk,
                                 IoCtl.IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,
                                 IntPtr.Zero,
                                 0,
                                 mm,
                                 (uint)mm.GetLength(),
                                 ref cbSize,
                                 IntPtr.Zero);

                        if (cbSize != 0)
                        {
                            devOut.Geometry = new DiskGeometry(mm.ToStruct<DISK_GEOMETRY_EX>());
                        }

                        // we're done!
                        FileIO.CloseHandle(disk);
                        disk = IntPtr.Zero;
                    }
                }

                // Free the mighty unmanaged resources ...
                mm.Free();

                // Add the new disk...
                disks.Add(devOut);
                devOut = null;

                // keep going.
                c++;
            }

            // Close the enumeration...
            Setup.SetupDiDestroyDeviceInfoList(hDev);

            // return the new disk array of all the disk device goodness ...
            return disks.ToArray();
        }

        [DllImport("kernel32.dll")]
        internal static extern uint GetLastError();
    }

    /// <summary>
    /// Describes the physical geometry of a disk.
    /// </summary>
    [TypeConverter(typeof(ExpandableObjectConverter))]
    public sealed class DiskGeometry
    {
        private DISK_GEOMETRY_EX info;

        internal DiskGeometry(DISK_GEOMETRY_EX g)
        {
            info = g;
        }

        /// <summary>
        /// Returns the capacity of the disk, in bytes.
        /// </summary>
        public ulong DiskSize
        {
            get
            {
                return info.DiskSize;
            }
        }

        /// <summary>
        /// Returns the number of cylinders on the disk.
        /// </summary>
        public long Cylinders
        {
            get
            {
                return info.Geometry.Cylinders;
            }
        }

        /// <summary>
        /// Returns the tracks per cylinder.
        /// </summary>
        public uint TracksPerCylinder
        {
            get
            {
                return info.Geometry.TracksPerCylinder;
            }
        }

        /// <summary>
        /// Returns the sectors per track.
        /// </summary>
        public uint SectorsPerTrack
        {
            get
            {
                return info.Geometry.SectorsPerTrack;
            }
        }

        /// <summary>
        /// Returns the bytes per sector.
        /// </summary>
        public uint BytesPerSector
        {
            get
            {
                return info.Geometry.BytesPerSector;
            }
        }

        /// <summary>
        /// Returns the total number of sectors on the disk.
        /// </summary>
        public long TotalSectors
        {
            get
            {
                return info.Geometry.TotalSectors;
            }
        }

        /// <summary>
        /// Returns the media type of the physical disk.
        /// </summary>
        public MediaType MediaType
        {
            get
            {
                return info.Geometry.MediaType;
            }
        }

        /// <summary>
        /// Converts this object into its string representation.
        /// </summary>
        /// <returns></returns>
        public override string ToString()
        {
            return info.Geometry.ToString();
        }
    }

    /// <summary>
    /// Device removal policies.
    /// </summary>
    public enum DeviceRemovalPolicy : uint
    {
        /// <summary>
        /// The device is not expected to be removed.
        /// </summary>
        /// <remarks></remarks>
        ExpectNoRemoval = 1,

        /// <summary>
        /// The device can be expected to be removed in an orderly fashion.
        /// </summary>
        /// <remarks></remarks>
        ExpectOrderlyRemoval = 2,

        /// <summary>
        /// The device can be expected to be removed without any preparation for removal.
        /// </summary>
        /// <remarks></remarks>
        ExpectSurpriseRemoval = 3
    }

    /// <summary>
    /// Disk Information
    /// </summary>
    public class DiskInfo
    {
        /// <summary>
        /// The physical device path.  This can be used as an alternate path for CreateFile() to open the physical disk.
        /// </summary>
        public string DevicePath { get; internal set; }

        /// <summary>
        /// The "friendly" device name.
        /// </summary>
        public string FriendlyName { get; internal set; }

        /// <summary>
        /// This is the InstanceId.  This can be used to raise the device properties dialog.
        /// </summary>
        public string InstanceId { get; internal set; }

        /// <summary>
        /// Physical location paths.
        /// </summary>
        public string[] LocationPaths { get; internal set; }

        /// <summary>
        /// Physical Hardware Ids.
        /// </summary>
        public string[] HardwareIds { get; internal set; }

        /// <summary>
        /// Specific location information.
        /// </summary>
        public string LocationInfo { get; internal set; }

        /// <summary>
        /// The UINumber
        /// </summary>
        public int UINumber { get; internal set; }

        /// <summary>
        /// The device description.
        /// </summary>
        public string Description { get; internal set; }

        /// <summary>
        /// The device's installation date.
        /// </summary>
        public DateTime InstallDate { get; internal set; }

        /// <summary>
        /// The PDO Name.
        /// </summary>
        public string PDOName { get; internal set; }

        /// <summary>
        /// The manufacturer.
        /// </summary>
        public string Manufacturer { get; internal set; }

        /// <summary>
        /// The device interface class Guid.
        /// </summary>
        public Guid DeviceInterfaceClassGuid { get; internal set; }

        /// <summary>
        /// The device interface Guid.
        /// </summary>
        public Guid DeviceClassGuid { get; internal set; }

        /// <summary>
        /// The InstanceId of the parent device.
        /// </summary>
        public string Parent { get; internal set; }

        /// <summary>
        /// The removal policy of the device.
        /// </summary>
        public DeviceRemovalPolicy RemovalPolicy { get; internal set; }

        /// <summary>
        /// Safe removal requirement.
        /// </summary>
        public bool SafeRemovalRequired { get; internal set; }

        /// <summary>
        /// The physical drive number.
        /// </summary>
        public uint PhysicalDevice { get; internal set; }

        /// <summary>
        /// Detailed disk geometry information.
        /// </summary>
        public DiskGeometry Geometry { get; internal set; }

        /// <summary>
        /// Converts this object into its string representation.
        /// </summary>
        /// <returns></returns>
        public override string ToString()
        {
            return FriendlyName + " [" + Geometry.ToString() + "]";
        }

        /// <summary>
        /// Opens the hardware properties dialog box for the specified instance id.
        /// </summary>
        /// <remarks></remarks>
        public void OpenDevicePropertiesDialog()
        {
            OpenDeviceProperties(InstanceId);
        }

        /// <summary>
        /// Opens the hardware properties dialog box for the specified instance id.
        /// </summary>
        /// <param name="hwnd">The owner window handle.</param>
        /// <remarks></remarks>
        public void OpenDevicePropertiesDialog(IntPtr hwnd)
        {
            OpenDeviceProperties(InstanceId, hwnd);
        }

        /// <summary>
        /// Opens the hardware properties dialog box for the specified instance id.
        /// </summary>
        /// <param name="InstanceId">The device InstanceId to use to raise the dialog box.</param>
        /// <param name="hwnd">The owner window handle.</param>
        /// <remarks></remarks>
        public static void OpenDeviceProperties(string InstanceId, IntPtr hwnd)
        {
            Util.DeviceProperties_RunDLL(hwnd, System.Diagnostics.Process.GetCurrentProcess().Handle, "/DeviceId \"" + InstanceId + "\"", 1);
        }

        /// <summary>
        /// Opens the hardware properties dialog box for the specified instance id.
        /// </summary>
        /// <param name="InstanceId">The device InstanceId to use to raise the dialog box.</param>
        /// <remarks></remarks>
        public static void OpenDeviceProperties(string InstanceId)
        {
            Util.DeviceProperties_RunDLL(IntPtr.Zero, System.Diagnostics.Process.GetCurrentProcess().Handle, "/DeviceId \"" + InstanceId + "\"", 1);
        }
    }
}